package com.whz.gateway.common.web;

import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;

class LocalServerHttpRequestDecorator extends ServerHttpRequestDecorator {
    private final HttpHeaders headers;
    private final CachedHttpOutputMessage outputMessage;

    LocalServerHttpRequestDecorator(ServerWebExchange exchange,
                                    HttpHeaders headers,
                                    CachedHttpOutputMessage outputMessage) {
        super(exchange.getRequest());
        this.headers = new HttpHeaders(headers);
        this.outputMessage = outputMessage;
    }

    @Override
    public HttpHeaders getHeaders() {
        long contentLength = headers.getContentLength();
        HttpHeaders httpHeaders = new HttpHeaders();
        httpHeaders.putAll(headers);
        if (contentLength > 0) {
            httpHeaders.setContentLength(contentLength);
        } else {
            httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
        }
        return httpHeaders;
    }

    @Override
    public Flux<DataBuffer> getBody() {
        return outputMessage.getBody();
    }
}